[bugfix] avoid attention padding tokens computation in pcg by Chen-0210 · Pull Request #17706 · sgl-project/sglang

Chen-0210 · 2026-01-25T10:59:52Z

Motivation

When PCG is enabled, the attention metadata is initialized with real_num_tokens, but the input tensor still contains padded tokens. Attention backends such as FlashInfer can not handle this well, which can lead to undefined behavior, including nan value, corrupted outputs (repeated !!!!!), result in abnormally long output lengths.

To fix this, exclude the padded tokens and make PCG more robust.

Modifications

Reverts PR fix: piecewise_cuda_graph get correct qo_indptr #21452
Slices tensors to real_num_tokens before/after the attention forward call — in unified_attention_with_output and unified_linear_attention_with_output, query/key/value and out_cache_loc are narrowed to exclude padding tokens before calling the backend, then restored out_cache_loc afterward, so attention only ever sees real tokens.

Accuracy Tests

Tested in CI

Benchmarking and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

Review Process

Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
After green CI and required approvals, ask Merge Oncalls to merge.

gemini-code-assist · 2026-01-25T11:00:13Z

Summary of Changes

Hello @Chen-0210, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request implements a bug fix to prevent padding tokens from being included in attention calculations within the Piecewise CUDA Graph (PCG) execution. By introducing a real_num_tokens field to ForwardBatch and modifying attention layers to process only the actual tokens, the change ensures more accurate and efficient attention computations, eliminating the influence of padding on model outputs.

Highlights

Padding Token Exclusion: Attention mechanisms in radix_attention.py and qwen3_next.py now explicitly ignore padding tokens by slicing input tensors (query, key, value, hidden_states, rope, sinks) based on a new real_num_tokens field, ensuring computations only involve actual data.
Introduction of real_num_tokens: A new optional integer field, real_num_tokens, has been added to the ForwardBatch class to accurately track the number of non-padding tokens in a batch.
Piecewise CUDA Graph Runner Refactoring: The piecewise_cuda_graph_runner.py has been updated to correctly pass and utilize real_num_tokens during graph warmup and capture. It also streamlines the handling of out_cache_loc and out_cache_loc_swa by making them dynamically created or passed via ForwardBatch rather than persistent graph runner attributes.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a bugfix for handling padding tokens within the Piecewise CUDA Graph (PCG) execution path. The core change is the addition of a real_num_tokens field to the ForwardBatch dataclass, which allows distinguishing actual tokens from padding. This field is then utilized in custom attention operations (unified_attention_with_output and gdn_with_output) to correctly slice input tensors, ensuring that only real tokens are processed during attention computation. Consequently, the logic in PiecewiseCudaGraphRunner for handling out_cache_loc has been simplified by removing pre-allocated tensors and passing them directly from the forward_batch.

The changes appear correct and effectively address the padding issue in PCG mode. I have one minor suggestion to improve code clarity by correcting a duplicated comment.

Oasis-Git

Since these changes affect a broad range of functionality, the full unit test suite should be carefully validated before approving this PR.

Chen-0210 · 2026-04-08T13:25:50Z

/rerun-failed-ci

Chen-0210 · 2026-04-08T16:48:46Z

/rerun-failed-ci

Chen-0210 · 2026-04-09T06:13:02Z

/rerun-failed-ci

Chen-0210 · 2026-04-09T07:21:05Z

/rerun-failed-ci

…-qwen3-next-v3

Chen-0210 · 2026-04-09T08:55:18Z

/rerun-failed-ci

Chen-0210 · 2026-04-10T02:50:38Z

/rerun-failed-ci

Chen-0210 · 2026-04-10T05:12:28Z

/rerun-failed-ci

Chen-0210 · 2026-04-10T08:23:02Z

/rerun-failed-ci

…-qwen3-next-v3

Chen-0210 · 2026-04-12T03:14:24Z

/rerun-failed-ci

Chen-0210 · 2026-04-12T05:11:17Z

/rerun-failed-ci

Chen-0210 · 2026-04-12T06:14:18Z

/rerun-failed-ci

Chen-0210 · 2026-04-12T09:56:24Z

/rerun-failed-ci

Chen-0210 · 2026-04-12T14:05:06Z

/rerun-failed-ci

Oasis-Git

I think this is a clean change. approved

ispobock · 2026-04-14T06:36:16Z

@Chen-0210 could you fix the test conflict?

Chen-0210 · 2026-04-14T07:45:10Z

@Chen-0210 could you fix the test conflict?

Done

…17706)" This reverts commit 6760c79.

init

3bedab6

Chen-0210 requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, Ying1123, ch-wan, hebiao064, hnyls2002, ispobock and merrymercy as code owners January 25, 2026 10:59

gemini-code-assist bot reviewed Jan 25, 2026

View reviewed changes

Comment thread python/sglang/srt/model_executor/forward_batch_info.py Outdated

Chen-0210 mentioned this pull request Jan 25, 2026

[Bug] Output token count abnormally increases with Qwen3-Next when --enable-piecewise-cuda-graph is set #17330

Closed

5 tasks

Chen-0210 changed the title ~~[bugfix] ignore padding tokens for attention in PCG~~ [bugfix] avoid attention padding tokens computation in pcg Jan 25, 2026

Chen-0210 added 3 commits January 26, 2026 01:35

clean

04b91bc

clean

0e96e36

clean

644da7a

zminglei mentioned this pull request Jan 27, 2026

[Perf] refactor piecewise cuda graph support of Qwen3-Next #17613

Merged

5 tasks

hebiao064 reviewed Jan 29, 2026

View reviewed changes

Comment thread python/sglang/srt/models/qwen3_next.py Outdated

hebiao064 mentioned this pull request Jan 29, 2026

[Tracking] refactor hybrid linear attention #17447

Open

6 tasks

ch-wan reviewed Jan 29, 2026

View reviewed changes

Comment thread python/sglang/srt/model_executor/forward_batch_info.py Outdated

Oasis-Git reviewed Jan 29, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/radix_attention.py Outdated

Chen-0210 requested review from CatherineSue, iforgetmyname, key4ng, ping1jing2, slin1237, yizhang2077 and zhyncs as code owners January 30, 2026 10:02

Chen-0210 added 3 commits April 9, 2026 11:42

Simplify piecewise cache loc replay

68fc842

Drop piecewise cache loc branching and test overrides

ff22072

Merge branch 'main' into fix-qwen3-next

1107938

Chen-0210 added 4 commits April 9, 2026 16:14

Preserve forward batch identity in unified attention

71db919

Merge remote-tracking branch 'chen0210/fix-qwen3-next' into codex/fix…

d5486e4

…-qwen3-next-v3

Simplify unified attention cache loc restore

8563107

Merge branch 'main' into fix-qwen3-next

2f84232

Merge branch 'main' into fix-qwen3-next

9427dd7

Chen-0210 added 2 commits April 10, 2026 16:59

Slice SWA cache loc in unified attention

00c2192

Merge remote-tracking branch 'chen0210/fix-qwen3-next' into codex/fix…

c2dd538

…-qwen3-next-v3

Oasis-Git approved these changes Apr 14, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into fix-qwen3-next

53bec38

ispobock merged commit 6760c79 into sgl-project:main Apr 14, 2026
35 of 74 checks passed

Chen-0210 deleted the fix-qwen3-next branch April 14, 2026 08:29

bingxche added a commit that referenced this pull request Apr 15, 2026

Revert "[bugfix] avoid attention padding tokens computation in pcg (#…

42d91b4

…17706)" This reverts commit 6760c79.

Conversation

Chen-0210 commented Jan 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist bot commented Jan 25, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Oasis-Git left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Chen-0210 commented Apr 8, 2026

Uh oh!

Chen-0210 commented Apr 8, 2026

Uh oh!

Chen-0210 commented Apr 9, 2026

Uh oh!

Chen-0210 commented Apr 9, 2026

Uh oh!

Chen-0210 commented Apr 9, 2026

Uh oh!

Chen-0210 commented Apr 10, 2026

Uh oh!

Chen-0210 commented Apr 10, 2026

Uh oh!

Chen-0210 commented Apr 10, 2026

Uh oh!

Chen-0210 commented Apr 12, 2026

Uh oh!

Chen-0210 commented Apr 12, 2026

Uh oh!

Chen-0210 commented Apr 12, 2026

Uh oh!

Chen-0210 commented Apr 12, 2026

Uh oh!

Chen-0210 commented Apr 12, 2026

Uh oh!

Oasis-Git left a comment

Choose a reason for hiding this comment

Uh oh!

ispobock commented Apr 14, 2026

Uh oh!

Chen-0210 commented Apr 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Chen-0210 commented Jan 25, 2026 •

edited

Loading